package utils

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
)

func generateKeyPair(bits int) (*rsa.PrivateKey, *rsa.PublicKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, &privateKey.PublicKey, nil
}

func exportRSAPrivateKeyAsPEM(privateKey *rsa.PrivateKey) string {
	privKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	privKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: privKeyBytes,
	})
	return string(privKeyPEM)
}

func exportRSAPublicKeyAsPEM(publicKey *rsa.PublicKey) (string, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		return "", err
	}
	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: pubKeyBytes,
	})
	return string(pubKeyPEM), nil
}

func GenerateStringKey() {
	privateKey, publicKey, err := generateKeyPair(2048)
	if err != nil {
		fmt.Fprintf(os.Stderr, "无法生成密钥对: %v\n", err)
		os.Exit(1)
	}

	// 将私钥转换为PEM格式
	privateKeyPEM := exportRSAPrivateKeyAsPEM(privateKey)
	fmt.Printf("私钥:\n%s\n", privateKeyPEM)

	// 将公钥转换为PEM格式
	publicKeyPEM, err := exportRSAPublicKeyAsPEM(publicKey)
	if err != nil {
		fmt.Fprintf(os.Stderr, "无法导出公钥: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("公钥:\n%s\n", publicKeyPEM)
}
